{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import os\n",
    "from PIL import Image\n",
    "from PIL import ImageFile\n",
    "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
    "IMAGE_SIZE = 224\n",
    "train_batch_size = 12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_labels = pd.read_csv(\"./data/labels.csv\")\n",
    "labels_str   = pd.read_csv(\"./data/sample_submission.csv\").drop(\"id\", axis = 1).columns.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels_str_lenth=len(labels_str)\n",
    "labels_map={}\n",
    "for index,labels_index in enumerate(labels_str):\n",
    "    labels_map[labels_index]=[index]\n",
    "train_labels['breed']=train_labels['breed'].map(labels_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[23]\n",
      "23\n"
     ]
    }
   ],
   "source": [
    "print(train_labels[train_labels[\"id\"]=='fa2a33c1dc8b39ad51738408b289a0de']['breed'].ravel().reshape(-1)[0])\n",
    "print(train_labels[train_labels[\"id\"]=='fa2a33c1dc8b39ad51738408b289a0de']['breed'].ravel().reshape(-1)[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the ./process_workspace/TrainSet.tfrecords exist, no need to run the  convert_to_tfrecord\n"
     ]
    }
   ],
   "source": [
    "def convert_to_tfrecord(rootpath, target_record_dir, target_record_filename):\n",
    "    target_tfrecord_file = target_record_dir+\"/\"+target_record_filename+\".tfrecords\"\n",
    "    os.makedirs(target_record_dir, mode=0o775, exist_ok=True)\n",
    "    if os.path.isfile(target_tfrecord_file) is True:\n",
    "        print(\"the \"+target_tfrecord_file+\" exist, no need to run the  convert_to_tfrecord\")\n",
    "        return\n",
    "    writer = tf.python_io.TFRecordWriter(target_tfrecord_file)\n",
    "    filenames = os.listdir(rootpath)\n",
    "    for name in filenames:\n",
    "        print(\"Processing image:\" + name)\n",
    "        img = Image.open(rootpath+\"/\"+name)\n",
    "        label = name.rstrip(\".jpg\")\n",
    "        \n",
    "        if img.mode == \"RGB\":\n",
    "            img = img.resize((IMAGE_SIZE, IMAGE_SIZE), Image.LANCZOS)\n",
    "            img_raw = img.tobytes()\n",
    "            example = tf.train.Example(\n",
    "                features=tf.train.Features(feature={\n",
    "                    \"label\": tf.train.Feature(int64_list=\n",
    "                                              tf.train.Int64List(value=train_labels[train_labels[\"id\"]==label]['breed'].ravel().reshape(-1)[0])), \n",
    "                    \"img_raw\": tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))}))\n",
    "            writer.write(example.SerializeToString())\n",
    "    writer.close()\n",
    "convert_to_tfrecord(\"./data/train\", \"./process_workspace\", \"TrainSet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_op(input_op, name, kh, kw, n_out, dh, dw, p):\n",
    "    n_in = input_op.get_shape()[-1].value\n",
    "\n",
    "    with tf.name_scope(name) as scope:\n",
    "        kernel = tf.get_variable(scope+\"w\",\n",
    "                                 shape=[kh, kw, n_in, n_out], dtype=tf.float32,\n",
    "                                 initializer=tf.contrib.layers.xavier_initializer_conv2d())\n",
    "        conv = tf.nn.conv2d(input_op, kernel, (1, dh, dw, 1), padding=\"SAME\")\n",
    "        bias_init_val = tf.constant(0.0, shape=[n_out], dtype=tf.float32)\n",
    "        biases = tf.Variable(bias_init_val, trainable=True, name='b')\n",
    "        z = tf.nn.bias_add(conv, biases)\n",
    "        # add L2 loss\n",
    "        weight_loss = tf.multiply(tf.nn.l2_loss(kernel), 0.001, name=\"weight_loss\")\n",
    "        weight_loss = tf.reduce_mean(weight_loss)\n",
    "        tf.add_to_collection('losses', weight_loss)\n",
    "        activation = tf.nn.relu(z, name=scope)\n",
    "        p += [kernel, biases]\n",
    "        return activation\n",
    "\n",
    "def fc_op(input_op, name, n_out, p, wl2=0.004, relu_tag=True):\n",
    "    n_in = input_op.get_shape()[-1].value\n",
    "\n",
    "    with tf.name_scope(name) as scope:\n",
    "        kernel = tf.get_variable(scope+\"w\", shape=[n_in, n_out], dtype=tf.float32,\n",
    "                                 initializer=tf.contrib.layers.xavier_initializer())\n",
    "        if relu_tag is True:\n",
    "            biases = tf.Variable(tf.constant(0.001, shape=[n_out], dtype=tf.float32), name=\"b\")\n",
    "            activation = tf.nn.relu_layer(input_op, kernel, biases, name=scope)\n",
    "        else:\n",
    "            biases = tf.Variable(tf.constant(0.001, shape=[n_out], dtype=tf.float32), name=\"b\")\n",
    "            activation = tf.matmul(input_op, kernel) + biases\n",
    "        # add L2 loss\n",
    "        weight_loss = tf.multiply(tf.nn.l2_loss(kernel), wl2, name=\"weight_loss\")\n",
    "        weight_loss = tf.reduce_mean(weight_loss)\n",
    "        tf.add_to_collection('losses', weight_loss)\n",
    "        p += [kernel, biases]\n",
    "        return activation\n",
    "\n",
    "def mpool_op(input_op, name, kh, kw, dh, dw):\n",
    "    return tf.nn.max_pool(input_op,\n",
    "                          ksize=[1, kh, kw, 1],\n",
    "                          strides=[1, dh, dw, 1],\n",
    "                          padding=\"SAME\",\n",
    "                          name=name)\n",
    "\n",
    "def inference_op(input_op, keep_prob):\n",
    "    p = []\n",
    "    conv1_1 = conv_op(input_op, name=\"conv1_1\", kh=3, kw=3, n_out=64, dh=1, dw=1, p=p)\n",
    "    conv1_2 = conv_op(conv1_1, name=\"conv1_2\", kh=3, kw=3, n_out=64, dh=1, dw=1, p=p)\n",
    "    conv1_2_dropout = tf.nn.dropout(conv1_2, 1.0, name=\"conv1_2_dropout\")\n",
    "    pool1 = mpool_op(conv1_2_dropout, name=\"pool1\", kh=2, kw=2, dw=2, dh=2)\n",
    "\n",
    "    conv2_1 = conv_op(pool1, name=\"conv2_1\", kh=3, kw=3, n_out=128, dh=1, dw=1, p=p)\n",
    "    conv2_2 = conv_op(conv2_1, name=\"conv2_2\", kh=3, kw=3, n_out=128, dh=1, dw=1, p=p)\n",
    "    conv2_2_dropout = tf.nn.dropout(conv2_2, 1.0, name=\"conv2_2_dropout\")\n",
    "    pool2 = mpool_op(conv2_2_dropout, name=\"pool2\", kh=2, kw=2, dw=2, dh=2)\n",
    "\n",
    "    conv3_1 = conv_op(pool2, name=\"conv3_1\", kh=3, kw=3, n_out=256, dh=1, dw=1, p=p)\n",
    "    conv3_2 = conv_op(conv3_1, name=\"conv3_2\", kh=3, kw=3, n_out=256, dh=1, dw=1, p=p)\n",
    "    conv3_3 = conv_op(conv3_2, name=\"conv3_3\", kh=3, kw=3, n_out=256, dh=1, dw=1, p=p)\n",
    "    conv3_3_dropout = tf.nn.dropout(conv3_3, 1.0, name=\"conv3_3_dropout\")\n",
    "    pool3 = mpool_op(conv3_3_dropout, name=\"pool3\", kh=2, kw=2, dw=2, dh=2)\n",
    "\n",
    "    conv4_1 = conv_op(pool3, name=\"conv4_1\", kh=3, kw=3, n_out=512, dh=1, dw=1, p=p)\n",
    "    conv4_2 = conv_op(conv4_1, name=\"conv4_2\", kh=3, kw=3, n_out=512, dh=1, dw=1, p=p)\n",
    "    conv4_3 = conv_op(conv4_2, name=\"conv4_3\", kh=3, kw=3, n_out=512, dh=1, dw=1, p=p)\n",
    "    conv4_3_dropout = tf.nn.dropout(conv4_3, 1.0, name=\"conv4_3_dropout\")\n",
    "    pool4 = mpool_op(conv4_3_dropout, name=\"pool4\", kh=2, kw=2, dw=2, dh=2)\n",
    "\n",
    "    conv5_1 = conv_op(pool4, name=\"conv5_1\", kh=3, kw=3, n_out=512, dh=1, dw=1, p=p)\n",
    "    conv5_2 = conv_op(conv5_1, name=\"conv5_2\", kh=3, kw=3, n_out=512, dh=1, dw=1, p=p)\n",
    "    conv5_3 = conv_op(conv5_2, name=\"conv5_3\", kh=3, kw=3, n_out=512, dh=1, dw=1, p=p)\n",
    "    conv5_3_dropout = tf.nn.dropout(conv5_3, 1.0, name=\"conv5_3_dropout\")\n",
    "    pool5 = mpool_op(conv5_3_dropout, name=\"pool5\", kh=2, kw=2, dw=2, dh=2)\n",
    "\n",
    "    shp = pool5.get_shape()\n",
    "    flattened_shape = shp[1].value * shp[2].value * shp[3].value\n",
    "    resh1 = tf.reshape(pool5, [-1, flattened_shape], name=\"resh1\")\n",
    "\n",
    "    fc6 = fc_op(resh1, name=\"fc6\", n_out=4096, p=p, wl2=0.004)\n",
    "    fc6_drop = tf.nn.dropout(fc6, keep_prob, name=\"fc6_drop\")\n",
    "\n",
    "    fc7 = fc_op(fc6_drop, name=\"fc7\", n_out=4096, p=p, wl2=0.004)\n",
    "    fc7_drop = tf.nn.dropout(fc7, keep_prob, name=\"fc7_drop\")\n",
    "\n",
    "\n",
    "    fc8 = fc_op(fc7_drop, name=\"fc8\", n_out=120, p=p)\n",
    "    softmax = tf.nn.softmax(fc8)\n",
    "    prediction = tf.argmax(softmax, 1)\n",
    "    return prediction, softmax, fc8, p\n",
    "\n",
    "\n",
    "def loss(logits, labels):\n",
    "    labels = tf.cast(labels, tf.int32)\n",
    "    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "        logits=logits, labels=labels, name='cross_entropy_per_example'\n",
    "    )\n",
    "    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')\n",
    "    tf.add_to_collection('losses', cross_entropy_mean)\n",
    "    return tf.add_n(tf.get_collection('losses'), name='total_loss')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_tfrecord(filenames, batch_num):\n",
    "    read_tfrecord_dataset = tf.contrib.data.TFRecordDataset(filenames)\n",
    "    read_tfrecord_dataset = read_tfrecord_dataset.map(_parse_function)\n",
    "    read_tfrecord_dataset = read_tfrecord_dataset.repeat(10)\n",
    "    read_tfrecord_dataset = read_tfrecord_dataset.batch(batch_num)\n",
    "    return read_tfrecord_dataset\n",
    "\n",
    "def save_model(sess, variable_list, checkpointfile):\n",
    "    saverdict = {}\n",
    "    for item in variable_list:\n",
    "        saverdict.update({item.name[:-2]: item})\n",
    "    saver = tf.train.Saver(saverdict)\n",
    "    save_path = saver.save(sess, checkpointfile)\n",
    "    print(\"model saved in the \"+save_path)\n",
    "def _parse_function(filenames):\n",
    "    features = {\"img_raw\": tf.FixedLenFeature([], tf.string),\n",
    "                \"label\": tf.FixedLenFeature(1, tf.int64)}\n",
    "    parsed_features = tf.parse_single_example(filenames, features)\n",
    "    image = tf.decode_raw(parsed_features['img_raw'], tf.uint8)\n",
    "    return image, parsed_features[\"label\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 1)\n"
     ]
    }
   ],
   "source": [
    "# DEBUG dock\n",
    "dataset_debug = read_tfrecord([\"./process_workspace/TrainSet.tfrecords\"], 2)\n",
    "with tf.Session() as sess:\n",
    "    iterator = dataset_debug.make_initializable_iterator()\n",
    "    next_element = iterator.get_next()\n",
    "    sess.run(iterator.initializer)\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    preprocess_batch_img_test = sess.run(next_element)\n",
    "    print(preprocess_batch_img_test[1].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset = read_tfrecord([\"./process_workspace/TrainSet.tfrecords\"], train_batch_size)\n",
    "keep_prod = tf.placeholder(tf.float32)\n",
    "image_holder = tf.placeholder(tf.float32, [train_batch_size, IMAGE_SIZE, IMAGE_SIZE, 3])\n",
    "label_holder = tf.placeholder(tf.int64, [train_batch_size])\n",
    "testimage_holder = tf.placeholder(tf.float32, [1, IMAGE_SIZE, IMAGE_SIZE, 3])\n",
    "# softmax output use\n",
    "prediction, softmax, fc8, p = inference_op(image_holder, keep_prod)\n",
    "loss = loss(fc8, label_holder)\n",
    "train_op = tf.train.AdamOptimizer(1e-5).minimize(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "########## train  0 step(s) start: #######\n"
     ]
    }
   ],
   "source": [
    "with tf.Session() as sess:\n",
    "    iterator = dataset.make_initializable_iterator()\n",
    "    next_element = iterator.get_next()\n",
    "    sess.run(iterator.initializer)\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    saverdict = {}\n",
    "    for item in p:\n",
    "        saverdict.update({item.name[:-2]: item})\n",
    "    saver = tf.train.Saver(saverdict)\n",
    "    if os.path.isfile(\"./process_workspace/DogBreedIdentificationVgg16_model.ckpt.meta\") is True:\n",
    "        print(\"Found checkpoint files , start to training from restored data\")\n",
    "        saver.restore(sess, \"./process_workspace/DogBreedIdentificationVgg16_model.ckpt\")\n",
    "        print(p[-1], sess.run(p[-1]))\n",
    "        print(\"Model restored.\\n\")\n",
    "    Step_Continue = 0\n",
    "    step = -1\n",
    "    # end in 800\n",
    "    for step in range(10):\n",
    "        if step < Step_Continue:\n",
    "            sess.run(next_element)\n",
    "            continue\n",
    "        print(\"########## train  %d step(s) start: #######\" % step, end=\"\\n\")\n",
    "        try:\n",
    "            preprocess_batch_img = sess.run(next_element)\n",
    "            _,  result_loss = sess.run([train_op, loss], feed_dict={\n",
    "                image_holder: preprocess_batch_img[0].reshape(-1, IMAGE_SIZE, IMAGE_SIZE, 3)/255 - 0.5,\n",
    "                label_holder: preprocess_batch_img[1].reshape(-1),\n",
    "                keep_prod: 0.5\n",
    "            })\n",
    "            print(\"result_loss:\")\n",
    "            print(result_loss)\n",
    "        except tf.errors.OutOfRangeError:\n",
    "            save_model(sess, p, \"./process_workspace/DogBreedIdentificationVgg16_model.ckpt\")\n",
    "            break\n",
    "        if step % 5 == 0:\n",
    "            save_model(sess, p, \"./process_workspace/DogBreedIdentificationVgg16_model.ckpt\")\n",
    "    save_model(sess, p, \"./process_workspace/CatsDogsVgg16_model.ckpt\")\n",
    "    print(\"########## train end at  %d  range: #######\" % (step+1), end=\"\\n\")\n",
    "    print(sess.run(p[-1]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
